import argparse
import numpy as np
import torch
import os

import mmcv
from mmcv import Config
from mmcv.parallel import collate, scatter

from mmdet.models import build_detector
from mmdet.datasets import replace_ImageToTensor
from mmdet.datasets.pipelines import Compose
from mmdet.core import bbox2result

from mmdet.datasets import build_dataloader, build_dataset, replace_ImageToTensor

from magicmind.python.runtime import *

from tqdm import tqdm


def parse_args():
    parser = argparse.ArgumentParser(description="MMDet test (and eval) a model")
    parser.add_argument("config", help="test config file path")
    parser.add_argument("mm_model", help="magicmind model, include data and graph")
    parser.add_argument("out_file", help="output file path")
    args = parser.parse_args()
    return args


def convert_outputs(outputs):
    return [torch.from_numpy(o.asnumpy()) for o in outputs]


def get_bbox(net_out, data, model):
    bbox_list = model.bbox_head.get_bboxes(
        net_out[:3], net_out[3:], data["img_metas"][0].data[0], rescale=True
    )
    bbox_results = [
        bbox2result(det_bboxes, det_labels, model.bbox_head.num_classes)
        for det_bboxes, det_labels in bbox_list
    ]
    return bbox_results


def main():
    args = parse_args()

    cfg = Config.fromfile(args.config)
    cfg.model.pretrained = None
    cfg.data.test.test_mode = True

    pt_model = build_detector(cfg.model, test_cfg=cfg.get("test_cfg"))

    # build model
    model = Model()
    model.deserialize_from_file(
        os.path.join(args.mm_model, "graph"), os.path.join(args.mm_model, "data")
    )
    engine = model.create_i_engine()
    context = engine.create_i_context()

    inputs = context.create_inputs()
    inputs[0].from_numpy(torch.rand(1, 720, 1280, 3).numpy())
    dev = Device()
    dev.id = 0
    dev.active()
    inputs[0].to(dev)
    outputs = context.create_outputs(inputs)

    queue = dev.create_queue()
    context.enqueue(inputs, outputs, queue)
    outputs_tensor = convert_outputs(outputs)

    dataset = build_dataset(cfg.data.test)
    data_loader = build_dataloader(
        dataset, samples_per_gpu=1, workers_per_gpu=1, dist=False, shuffle=False
    )

    results = []
    for i, info in tqdm(enumerate(data_loader)):
        image = info["img"][0].permute(0, 2, 3, 1)
        inputs[0].from_numpy(image.numpy())
        context.enqueue(inputs, outputs, queue)
        queue.sync()
        outputs_tensor = convert_outputs(outputs)
        bbox_list = get_bbox(outputs_tensor, info, pt_model)
        results.extend(bbox_list)

    mmcv.dump(results, args.out_file)
    print(f"save detection results to {args.out_file}")


if __name__ == "__main__":
    main()
